#include <libwebsockets.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <pthread.h>

#define UDP_PORT 4444
#define WEB_SOCKET_PORT 9001

static struct lws_context *context;
static struct lws_vhost *vhost;
static struct lws_protocols protocols[] = {
    {
        "udp-to-websocket",
        NULL, // 回调函数在下面定义
        0,
        1024,
    },
    { NULL, NULL, 0, 0 } // 终止符
};

// 全局变量用于存储所有连接的WebSocket客户端
static struct lws *clients[100];
static int client_count = 0;
static pthread_mutex_t client_mutex = PTHREAD_MUTEX_INITIALIZER;

// WebSocket协议回调函数
static int callback_udp_to_websocket(struct lws *wsi, enum lws_callback_reasons reason, void *user, void *in, size_t len) {
    switch (reason) {
        case LWS_CALLBACK_ESTABLISHED:
            pthread_mutex_lock(&client_mutex);
            if (client_count < sizeof(clients)/sizeof(clients[0])) {
                clients[client_count++] = wsi;
                printf("新的WebSocket客户端连接，当前客户端数: %d\n", client_count);
            }
            pthread_mutex_unlock(&client_mutex);
            break;
            
        case LWS_CALLBACK_CLOSED:
            pthread_mutex_lock(&client_mutex);
            for (int i = 0; i < client_count; i++) {
                if (clients[i] == wsi) {
                    // 移除客户端
                    for (int j = i; j < client_count - 1; j++) {
                        clients[j] = clients[j+1];
                    }
                    client_count--;
                    printf("WebSocket客户端断开连接，当前客户端数: %d\n", client_count);
                    break;
                }
            }
            pthread_mutex_unlock(&client_mutex);
            break;
            
        default:
            break;
    }
    return 0;
}

// 初始化WebSocket协议回调函数
static void init_protocols() {
    protocols[0].callback = callback_udp_to_websocket;
}

// UDP监听线程
void *udp_listener_thread(void *arg) {
    int udp_socket;
    struct sockaddr_in server_addr, client_addr;
    socklen_t client_len = sizeof(client_addr);
    char buffer[1024];
    int recv_len;
    
    // 创建UDP socket
    if ((udp_socket = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
        perror("创建UDP socket失败");
        return NULL;
    }
    
    // 设置服务器地址
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = htonl(INADDR_ANY);
    server_addr.sin_port = htons(UDP_PORT);
    
    // 绑定socket
    if (bind(udp_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("绑定UDP socket失败");
        close(udp_socket);
        return NULL;
    }
    
    printf("开始在UDP端口 %d 监听...\n", UDP_PORT);
    
    while (1) {
        // 接收UDP数据
        recv_len = recvfrom(udp_socket, buffer, sizeof(buffer), 0, 
                           (struct sockaddr *)&client_addr, &client_len);
        if (recv_len < 0) {
            perror("接收UDP数据失败");
            continue;
        }
        
        printf("收到来自 %s:%d 的UDP数据，长度: %d\n", 
               inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port), recv_len);
        
        // 向所有WebSocket客户端发送数据
        pthread_mutex_lock(&client_mutex);
        if (client_count > 0) {
            for (int i = 0; i < client_count; i++) {
                // 分配内存并复制数据
                unsigned char *buf = malloc(LWS_PRE + recv_len);
                memcpy(buf + LWS_PRE, buffer, recv_len);
                
                // 发送数据
                lws_write(clients[i], buf + LWS_PRE, recv_len, LWS_WRITE_BINARY);
                
                free(buf);
            }
            printf("数据已发送给 %d 个WebSocket客户端\n", client_count);
        }
        pthread_mutex_unlock(&client_mutex);
    }
    
    close(udp_socket);
    return NULL;
}

int main(int argc, char *argv[]) {
    struct lws_context_creation_info info;
    pthread_t udp_thread;
    
    // 初始化协议回调函数
    init_protocols();
    
    // 设置上下文创建信息
    memset(&info, 0, sizeof(info));
    info.port = WEB_SOCKET_PORT;
    info.protocols = protocols;
    info.gid = -1;
    info.uid = -1;
    
    // 创建libwebsocket上下文
    context = lws_create_context(&info);
    if (!context) {
        fprintf(stderr, "创建libwebsocket上下文失败\n");
        return -1;
    }
    
    printf("WebSocket服务器启动在端口 %d\n", WEB_SOCKET_PORT);
    
    // 创建UDP监听线程
    if (pthread_create(&udp_thread, NULL, udp_listener_thread, NULL) != 0) {
        fprintf(stderr, "创建UDP监听线程失败\n");
        lws_context_destroy(context);
        return -1;
    }
    
    // 主循环处理WebSocket事件
    while (1) {
        lws_service(context, 50); // 50ms超时
    }
    
    // 清理资源
    lws_context_destroy(context);
    pthread_join(udp_thread, NULL);
    
    return 0;
}
